library(ggplot2)
library(reshape2)
library(scales)
library(grid)
library(gridExtra)
library(cowplot)
library(extrafont)

rm(list = ls())
setwd("")

#Traditional: Pre=0, Placebo=0, Multiple=0, Mode=T
#Proposed: Pre=1, Placebo=1, Multiple=1, Mode=O

# Constants
SIGMA2 <- 1
VSTAR <- .002
F <- 2 #Number of follow-up surveys

# Parameters to Vary
#A: Proportion of subjects attempted for treatment that are successfully treated
A <- .25
#T: Marginal cost of attempting treatment or placebo contact
T <- 3
#s: Marginal cost of survey, regardless of mode
s <- 5
#r_tel: response rate to follow-up survey
r_tel <- .07
#r_online: response rate to follow-up survey
r_online <- .75
#rho2: R-sq of regression outcome
rho2 <- .81

# Formulas
traditional <- function(A = .25, T = 3, s = 5, r_tel = .07, F = 2, rho2 = NA, r_online = NA){
  4 * (1/A^2) * (SIGMA2 / VSTAR) * (s*F + T / (2 * r_tel)) 
}
proposed <- function(A = .25, T = 3, s = 5, r_online = .75, rho2 = .81, F = 2, r_tel = NA) {
  4 * (1 - rho2) *
    (SIGMA2 / VSTAR) * (s*F + T / (A * r_online) + s / (A * r_online)) 
}

#A: Proportion of subjects attempted for treatment that are successfully treated
traditional.A <- function(A) traditional(A=A)
proposed.A <- function(A) proposed(A=A)
A.plot <- ggplot(data.frame(x=c(0.2,1)), aes(x)) + 
  stat_function(fun=traditional.A, aes(colour = "Traditional")) +
  stat_function(fun=proposed.A, aes(colour = "All Four Practices")) +
  ylab('Variable Cost') + xlab('Changing treatment application rate (A), all else fixed') +
  scale_colour_manual("", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))

#T: Marginal cost of attempting treatment or placebo contact
traditional.T <- function(T) traditional(T=T)
proposed.T <- function(T) proposed(T=T)
T.plot <- ggplot(data.frame(x=c(0,50)), aes(x)) + 
  stat_function(fun=traditional.T, aes(colour = "Traditional")) +
  stat_function(fun=proposed.T, aes(colour = "All Four Practices")) +
  ylab('Variable Cost') + xlab('Changing marginal treatment cost (T), all else fixed') +
  scale_colour_manual("", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))

#s: Marginal cost of survey, regardless of mode
traditional.s <- function(s) traditional(s=s)
proposed.s <- function(s) proposed(s=s)
s.plot <- ggplot(data.frame(x=c(0,50)), aes(x)) + 
  stat_function(fun=traditional.s, aes(colour = "Traditional")) +
  stat_function(fun=proposed.s, aes(colour = "All Four Practices")) +
  ylab('Variable Cost') + xlab('Changing marginal survey cost (s), all else fixed') +
  scale_colour_manual("", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))

#r_tel: response rate to follow-up survey, telephone
#r_online: response rate to follow-up survey, online
traditional.r_tel <- function(r_tel) traditional(r_tel=r_tel)
proposed.r_online <- function(r_online) proposed(r_online=r_online)
response.plot <- ggplot(data.frame(x=c(0.02,1)), aes(x)) + 
  stat_function(fun=traditional.r_tel, aes(colour = "Traditional")) +
  stat_function(fun=proposed.r_online, aes(colour = "All Four Practices")) +
  ylab('Variable Cost') + xlab('Changing follow-up survey response rate (R), all else fixed') +
  scale_colour_manual("", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))


#rho2: R-sq of regression outcome
traditional.rho2 <- function(rho2) traditional(rho2=rho2)
proposed.rho2 <- function(rho2) proposed(rho2=rho2)
rho2.plot <- ggplot(data.frame(x=c(0.02,1)), aes(x)) + 
  stat_function(fun=traditional.rho2, aes(colour = "Traditional")) +
  stat_function(fun=proposed.rho2, aes(colour = "All Four Practices")) +
  ylab('Variable Cost') + xlab(expression(paste('Changing ',rho^'2',', all else fixed'))) +
  scale_colour_manual("", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))

#F: Number of surveys
traditional.F <- function(F) traditional(F=F)
proposed.F <- function(F) proposed(F=F)
range.F <- seq(1,5,1)
traditional.F.cost <- sapply(range.F, traditional.F)
proposed.F.cost <- sapply(range.F, proposed.F)
f.data <- data.frame(cbind(c(traditional.F.cost, proposed.F.cost),
                           c(rep("Traditional Design", 5), rep("Design With All Four Practices", 5)), 
                           range.F), stringsAsFactors = FALSE)
f.data$V1 <- as.numeric(f.data$V1)
f.data$range.F <- as.numeric(f.data$range.F)
f.plot <- ggplot(f.data, aes(range.F, V1, fill = V2)) +
  geom_bar(stat = "identity", position="dodge") +
  ylab('Variable Cost') + xlab('Changing F, all else fixed') +
  scale_fill_manual(name = "", values = c("red", "blue")) +
  scale_y_continuous(labels = scales::dollar) + 
  theme(legend.position="bottom", text = element_text(family = 'Verdana'))

grid_arrange_shared_legend <- function(...) {
  plots <- list(...)
  g <- ggplotGrob(plots[[1]] + theme(legend.position="bottom"))$grobs
  legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
  lheight <- sum(legend$height)
  grid.arrange(
    do.call(arrangeGrob, lapply(plots, function(x)
      x + theme(legend.position="none"))),
    legend,
    ncol = 1,
    heights = unit.c(unit(1, "npc") - lheight, lheight))
}
#Figure 7: How Variable Costs Change with Parameters
tiff(filename = "output/changing_parameters.tiff", width = 1500, height = 800)#, res = 150)
  grid_arrange_shared_legend(rho2.plot, response.plot, s.plot, T.plot, A.plot, f.plot)
dev.off()
